/*	$Id: radix.c 6 2008-07-29 02:30:10Z phrakt $	*/
/*
 * Copyright (c) 2003 Jean-Francois Brousseau <jfb@openbsd.org>
 * All rights reserved.
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions
 * are met:
 *
 * 1. Redistributions of source code must retain the above copyright
 *    notice, this list of conditions and the following disclaimer.
 * 2. The name of the author may not be used to endorse or promote products
 *    derived from this software without specific prior written permission.
 *
 * THIS SOFTWARE IS PROVIDED ``AS IS'' AND ANY EXPRESS OR IMPLIED WARRANTIES,
 * INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY
 * AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL
 * THE AUTHOR BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
 * EXEMPLARY, OR CONSEQUENTIAL  DAMAGES (INCLUDING, BUT NOT LIMITED TO,
 * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS;
 * OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY,
 * WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR
 * OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF
 * ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 */

/*
* Radix trees
* ===========
*
* This code provides a simple generic API to manipulate radix trees (commonly
* used in routing tables, among other things).  A radix tree is a binary tree
* in which nodes are located depending on a key, which is an arbitrary string
* of bits.
*
*/

#include <sys/types.h>

#include <stdlib.h>
#include <string.h>

#include "radix.h"

/* node types */
#define RADIX_NODE_BRANCH 1
#define RADIX_NODE_LEAF   2



/* lookup types for pf_radix_findnode() */
#define RADIX_FIND_EXACT  1
#define RADIX_FIND_PARENT 2


typedef struct radix_node {
	int node_type;
	uint8_t *key;
	uint8_t *mask;

	struct radix_node *parent;

	union {
		struct {
			void   *data;
		} leaf;

		struct {
			struct radix_node *left;
			struct radix_node *right;
		} branch;
	} rn_un;
} RADIXNODE;

/* helpful macros */
#define rn_data   rn_un.leaf.data
#define rn_left   rn_un.branch.left
#define rn_right  rn_un.branch.right


struct radix_tree {
	int keylen;

	RADIXNODE *root;

	/* optional destructor called on the <data> part of nodes */
	void (*destructor) (void *);
};

/* test if the b'th bit of the key k is 1 or 0 */
#define RADIX_TESTBIT(k, b) \
	((k)[(b)/sizeof(uint8_t)] & ((b)%sizeof(uint8_t)))


static RADIXNODE* pf_radix_findnode  (RADIXTREE *, const uint8_t *, int);
static RADIXNODE* pf_radix_allocnode (int, const uint8_t *);
static void       pf_radix_freenode  (RADIXTREE *, RADIXNODE *);
static int        pf_radix_keycmp    (const uint8_t *, const uint8_t *, int);


/*
* pf_radix_create()
*
* Create a new radix tree and return a pointer to it.  The <klen> argument
* specifies the key length to use (in bits).  All other functions called on the
* returned tree and using a key argument expect that argument to be at least
* <klen> bits long.
* Returns a pointer to the newly created radix tree, or NULL on failure.
*/

RADIXTREE*
pf_radix_create(int klen, int flags, void (*destructor)(void *))
{
	struct radix_tree *tp;

	tp = (struct radix_tree *) calloc(1, sizeof(*tp));
	if (tp == NULL) {
		return (NULL);
	}

	tp->keylen = klen;
	tp->destructor = destructor;

#if 0
	tp->root = pf_radix_allocnode(RADIX_NODE_BRANCH);
#endif

	return (tp);
}


/*
* pf_radix_free()
*
* Free a radix tree and all associated data.
*/

void
pf_radix_free(void *vp)
{
	RADIXTREE *tree;

	tree = (RADIXTREE *)vp;

	if (tree != NULL) {
		if (tree->root != NULL) {
			pf_radix_freenode(tree, tree->root);
		}

		free(tree);
	}
}


/*
* pf_radix_insert()
*
* Insert a new node with key <key> in the tree <tree>, representing the
* data found at <data>.  Any branch nodes required are created at the same
* time.
* Returns 0 on success, or -1 on failure.
*/

int
pf_radix_insert(struct radix_tree *tree, const uint8_t *key, const void *data)
{
	int i;
	RADIXNODE *rn, *parent;

	if ((tree == NULL) || (key == NULL)) {
		return (-1);
	}

	parent = tree->root;
	for (i = 0; i < tree->keylen; i++) {
		if (RADIX_TESTBIT(key, i) == 1)
			parent = parent->rn_left;
		else
			parent = parent->rn_right;
	}

	rn = pf_radix_allocnode(RADIX_NODE_LEAF, NULL);
	if (rn == NULL) {
		return (-1);
	}

	rn->rn_data = data;

	return (0);
}


/*
* pf_radix_lookup()
*
* Lookup a node inside the radix tree <tree> whose key is <key>.  Returns the
* data pointer associated with that node if found, NULL otherwise.
*/

const void*
pf_radix_lookup(struct radix_tree *tree, const uint8_t *key)
{
	RADIXNODE *rn;

	if ((tree == NULL) || (key == NULL)) {
		return (NULL);
	}

	rn = pf_radix_findnode(tree, key, RADIX_FIND_EXACT);
	if (rn == NULL) {
		return (NULL);
	}

	return (rn->rn_data);
}


/*
* pf_radix_delete()
*
* Delete a node from the radix tree <tree> with key <key>.
*
*/

int
pf_radix_delete(RADIXTREE *tree, const uint8_t *key)
{
	RADIXNODE *rnp;

	rnp = pf_radix_findnode(tree, key, RADIX_FIND_EXACT);
	if (rnp == NULL) {
		return (-1);
	}

	if (rnp->parent != NULL) {
		if (rnp->parent->rn_left == rnp) {
			rnp->parent->rn_left = NULL;
		}
		else if (rnp->parent->rn_right == rnp) {
			rnp->parent->rn_right = NULL;
		}
	}

	pf_radix_freenode(tree, rnp);

	return (0);
}


/*
* pf_radix_findnode()
*
* This function is used to find a node inside the radix tree <tree>.  It can
* behave in two ways depending on the value of the <mode> argument.
* Returns a pointer to the radix node that matched, or NULL on failure.
*/

RADIXNODE*
pf_radix_findnode(RADIXTREE *tree, const uint8_t *key, int how)
{
	int bitoff;
	uint8_t mask;
	RADIXNODE *rn;

	rn = tree->root;

	for (bitoff = 0; bitoff < tree->keylen; bitoff++) {
		mask = 0x80 >> (bitoff % 8);
		if ( & mask) {
			if (rn->rn_left == NULL) {
				break;
			}

			rn = rn->rn_left;
		}
		else {
			if (rn->rn_right == NULL) {
				break;
			}
			rn = rn->rn_right;
		}
	}

	if ((how == RADIX_FIND_EXACT) && (bitoff != tree->keylen - 1)) {
		return (NULL);
	}

	return (rn);
}


/*
* pf_radix_allocnode()
*
* Allocate a new radix node, with type <type> and key <key>.  If no key is
* yet defined for the new node, the argument can be set to NULL.
* Returns a pointer to the newly allocated node on success, or NULL on failure.
*/

RADIXNODE*
pf_radix_allocnode(int type, const uint8_t *key)
{
	RADIXNODE *rn;

	rn = (RADIXNODE *) calloc(1, sizeof(*rn));
	if (rn == NULL) {
		return (NULL);
	}

	rn->node_type = type;

	return (rn);
}


/*
* pf_radix_freenode()
*
* Free a radix node and all subnodes.
*/

static void
pf_radix_freenode(RADIXTREE *tree, RADIXNODE *rnp)
{
	if (rnp->key != NULL)
		free(rnp->key);

	if (rnp->node_type == RADIX_NODE_LEAF) {
		if ((rnp->rn_data != NULL) && (tree->destructor != NULL)) {
			(*tree->destructor)(rnp->rn_data);
		}
	}
	else {
		if (rnp->rn_left != NULL) {
			pf_radix_freenode(tree, rnp->rn_left);
		}

		if (rnp->rn_right != NULL) {
			pf_radix_freenode(tree, rnp->rn_right);
		}
	}

	free(rnp);
}


/*
* pf_radix_keycmp()
*
* Compare the first <len> bits of the two radix keys <k1> and <k2>.
* Returns 0 if both keys are the same for the first <len> bits, greater than
* 0 if <k1> is greater than <k2>, and less than 0 otherwise.
*/

int
pf_radix_keycmp(const uint8_t *k1, const uint8_t *k2, int len)
{
	int ret;
	uint8_t k1bits, k2bits;
	int bytes, bits;

	bytes = len / 8;
	bits = len % 8;

	ret = memcmp(k1, k2, bytes);
	if (ret != 0) {
		/* no point in checking the remaining bits */
		return (ret);
	}

	k1bits = k1[bytes] & (0xff00 >> bits);
	k2bits = k2[bytes] & (0xff00 >> bits);

	if (k1bits > k2bits) {
		return (1);
	}
	else if (k1bits < k2bits) {
		return (-1);
	}

	return (0);
}
